﻿#-------------------------------------------------------------------------------
# Name:        Driver
# Purpose:     Windows driver manager
#
# Author:      SUPER_紫电  QQ:911344755
#
# Created:     15/08/2022
# Copyright:   (c) SUPER_紫电 2022 All rights reserved.
# Licence:     <Modified BSD licence>
#-------------------------------------------------------------------------------

import ctypes
import pywintypes
import win32con
import win32service
import winerror
from win32 import *

DRIVER_INSTALL = 1
DRIVER_REMOVE = 2
DRIVER_SYSTEM_INSTALL = 3
DRIVER_SYSTEM_UNINSTALL = 4

gHandle = None

def OpenDriver(DriverId):
    try:
        global gHandle
        gHandle = CreateFile(
            "\\\\.\\" + DriverId,
            win32con.GENERIC_READ | win32con.GENERIC_WRITE,
            0,
            None,
            win32con.OPEN_EXISTING,
            win32con.FILE_ATTRIBUTE_NORMAL,
            None
        )
        return gHandle != INVALID_HANDLE_VALUE
    except pywintypes.error as e:
        if e.winerror != winerror.ERROR_FILE_NOT_FOUND:
            raise e

def ManageDriver(DriverId, DriverPath, Function):
    global gHandle
    rCode = False
    if DriverId == "" or DriverPath == "":
        return False
    hSCManager = win32service.OpenSCManager(
        None,
        None,
        win32service.SC_MANAGER_ALL_ACCESS
    )
    if hSCManager == INVALID_HANDLE_VALUE:
        return False
    if Function == DRIVER_INSTALL:
        if InstallDriver(hSCManager, DriverId, DriverPath):
            rCode = StartDriver(hSCManager, DriverId)
    elif Function == DRIVER_REMOVE:
        if not IsSystemInstallDriver(hSCManager, DriverId, DriverPath):
            StopDriver(hSCManager, DriverId)
            rCode = RemoveDriver(hSCManager, DriverId)
    elif Function == DRIVER_SYSTEM_INSTALL:
        if IsSystemInstallDriver(hSCManager, DriverId, DriverPath):
            rCode = True
        else:
            if not OpenDriver(DriverId):
                StopDriver(hSCManager, DriverId)
                RemoveDriver(hSCManager, DriverId)
                if InstallDriver(hSCManager, DriverId, DriverPath):
                    StartDriver(hSCManager, DriverId)
                OpenDriver(DriverId)
            rCode = SystemInstallDriver(hSCManager, DriverId, DriverPath)
    elif Function == DRIVER_SYSTEM_UNINSTALL:
        if not IsSystemInstallDriver(hSCManager, DriverId, DriverPath):
            rCode = Ture
        else:
            if gHandle != INVALID_HANDLE_VALUE:
                CloseHandle(gHandle)
                gHandle = INVALID_HANDLE_VALUE
            if StopDriver(hSCManager, DriverId):
                rCode = RemoveDriver(hSCManager, DriverId)
    if hSCManager != INVALID_HANDLE_VALUE:
        win32service.CloseServiceHandle(hSCManager)
    return rCode

def InstallDriver(hSCManager, DriverId, DriverPath):
    try:
        hService = win32service.CreateService(
            hSCManager,
            DriverId,
            DriverId,
            win32service.SERVICE_ALL_ACCESS,
            win32service.SERVICE_KERNEL_DRIVER,
            win32service.SERVICE_DEMAND_START,
            win32service.SERVICE_ERROR_NORMAL,
            DriverPath,
            None,
            0,
            None,
            None,
            None
        )
        win32service.CloseServiceHandle(hService)
        return True
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_EXISTS:
            return True
        else:
            raise e

def SystemInstallDriver(hSCManager, DriverId, DriverPath):
    try:
        hService = win32service.OpenService(
            hSCManager,
            DriverId,
            win32service.SERVICE_ALL_ACCESS
        )
        if hService != INVALID_HANDLE_VALUE:
            rCode = win32service.ChangeServiceConfig(
                hService,
                win32service.SERVICE_KERNEL_DRIVER,
                win32service.SERVICE_AUTO_START,
                win32service.SERVICE_ERROR_NORMAL,
                DriverPath,
                None,
                None,
                None,
                None,
                None,
                None
            )
            win32service.CloseServiceHandle(hService)
        return rCode
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_DOES_NOT_EXIST:
            return True
        else:
            raise e

def RemoveDriver(hSCManager, DriverId):
    try:
        hService = win32service.OpenService(
            hSCManager,
            DriverId,
            win32service.SERVICE_ALL_ACCESS
        )
        pServiceConfig = win32service.QueryServiceConfig(hService)
        # If service is set to load automatically, don't delete it!
        if pServiceConfig[1] == win32service.SERVICE_DEMAND_START:
            rCode = win32service.DeleteService(hService)
        win32service.CloseServiceHandle(hService)
        return rCode
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_DOES_NOT_EXIST:
            return True
        else:
            raise e

def StartDriver(hSCManager, DriverId):
    try:
        hService = win32service.OpenService(
            hSCManager,
            DriverId,
            win32service.SERVICE_ALL_ACCESS
        )
        if hService != INVALID_HANDLE_VALUE:
            win32service.StartService(hService, None)
            win32service.CloseServiceHandle(hService)
            return True
        else:
            return False
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_ALREADY_RUNNING:
            return True
        else:
            raise e

def StopDriver(hSCManager, DriverId):
    try:
        hService = win32service.OpenService(
            hSCManager,
            DriverId,
            win32service.SERVICE_ALL_ACCESS
        )
        if hService != INVALID_HANDLE_VALUE:
            rCode = win32service.ControlService(
                hService,
                win32service.SERVICE_CONTROL_STOP
            )
            win32service.CloseServiceHandle(hService)
        return rCode
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_NOT_ACTIVE or \
            winerror.ERROR_SERVICE_DOES_NOT_EXIST:
            return True
        else:
            raise e

def IsSystemInstallDriver(hSCManager, DriverId, DriverPath):
    try:
        hService = win32service.OpenService(
            hSCManager,
            DriverId,
            win32service.SERVICE_ALL_ACCESS
        )
        if hService != INVALID_HANDLE_VALUE:
            lpServiceConfig = win32service.QueryServiceConfig(hService)
            rCode = lpServiceConfig[1] == win32service.SERVICE_AUTO_START
            win32service.CloseServiceHandle(hService)
        return rCode
    except pywintypes.error as e:
        if e.winerror == winerror.ERROR_SERVICE_DOES_NOT_EXIST:
            return True
        else:
            raise e
